{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "google",
   "metadata": {},
   "source": [
    "##### Copyright 2023 Google LLC."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "apache",
   "metadata": {},
   "source": [
    "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "you may not use this file except in compliance with the License.\n",
    "You may obtain a copy of the License at\n",
    "\n",
    "    http://www.apache.org/licenses/LICENSE-2.0\n",
    "\n",
    "Unless required by applicable law or agreed to in writing, software\n",
    "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "See the License for the specific language governing permissions and\n",
    "limitations under the License.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "basename",
   "metadata": {},
   "source": [
    "# nurses_cp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "link",
   "metadata": {},
   "source": [
    "<table align=\"left\">\n",
    "<td>\n",
    "<a href=\"https://colab.research.google.com/github/google/or-tools/blob/main/examples/notebook/contrib/nurses_cp.ipynb\"><img src=\"https://raw.githubusercontent.com/google/or-tools/main/tools/colab_32px.png\"/>Run in Google Colab</a>\n",
    "</td>\n",
    "<td>\n",
    "<a href=\"https://github.com/google/or-tools/blob/main/examples/contrib/nurses_cp.py\"><img src=\"https://raw.githubusercontent.com/google/or-tools/main/tools/github_32px.png\"/>View source on GitHub</a>\n",
    "</td>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "doc",
   "metadata": {},
   "source": [
    "First, you must install [ortools](https://pypi.org/project/ortools/) package in this colab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "install",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install ortools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "code",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "from ortools.constraint_solver import pywrapcp\n",
    "\n",
    "\n",
    "def main():\n",
    "  # Creates the solver.\n",
    "  solver = pywrapcp.Solver(\"schedule_shifts\")\n",
    "\n",
    "  num_nurses = 4\n",
    "  num_shifts = 4  # Nurse assigned to shift 0 means not working that day.\n",
    "  num_days = 7\n",
    "  # [START]\n",
    "  # Create shift variables.\n",
    "  shifts = {}\n",
    "\n",
    "  for j in range(num_nurses):\n",
    "    for i in range(num_days):\n",
    "      shifts[(j, i)] = solver.IntVar(0, num_shifts - 1,\n",
    "                                     \"shifts(%i,%i)\" % (j, i))\n",
    "  shifts_flat = [\n",
    "      shifts[(j, i)] for j in range(num_nurses) for i in range(num_days)\n",
    "  ]\n",
    "\n",
    "  # Create nurse variables.\n",
    "  nurses = {}\n",
    "\n",
    "  for j in range(num_shifts):\n",
    "    for i in range(num_days):\n",
    "      nurses[(j, i)] = solver.IntVar(0, num_nurses - 1,\n",
    "                                     \"shift%d day%d\" % (j, i))\n",
    "  # Set relationships between shifts and nurses.\n",
    "  for day in range(num_days):\n",
    "    nurses_for_day = [nurses[(j, day)] for j in range(num_shifts)]\n",
    "\n",
    "    for j in range(num_nurses):\n",
    "      s = shifts[(j, day)]\n",
    "      solver.Add(s.IndexOf(nurses_for_day) == j)\n",
    "  # Make assignments different on each day\n",
    "  for i in range(num_days):\n",
    "    solver.Add(solver.AllDifferent([shifts[(j, i)] for j in range(num_nurses)]))\n",
    "    solver.Add(solver.AllDifferent([nurses[(j, i)] for j in range(num_shifts)]))\n",
    "  # Each nurse works 5 or 6 days in a week.\n",
    "  for j in range(num_nurses):\n",
    "    solver.Add(solver.Sum([shifts[(j, i)] > 0 for i in range(num_days)]) >= 5)\n",
    "    solver.Add(solver.Sum([shifts[(j, i)] > 0 for i in range(num_days)]) <= 6)\n",
    "  # Create works_shift variables. works_shift[(i, j)] is True if nurse\n",
    "  # i works shift j at least once during the week.\n",
    "  works_shift = {}\n",
    "\n",
    "  for i in range(num_nurses):\n",
    "    for j in range(num_shifts):\n",
    "      works_shift[(i, j)] = solver.BoolVar(\"nurse%d shift%d\" % (i, j))\n",
    "\n",
    "  for i in range(num_nurses):\n",
    "    for j in range(num_shifts):\n",
    "      solver.Add(works_shift[(\n",
    "          i, j)] == solver.Max([shifts[(i, k)] == j for k in range(num_days)]))\n",
    "\n",
    "  # For each shift (other than 0), at most 2 nurses are assigned to that shift\n",
    "  # during the week.\n",
    "  for j in range(1, num_shifts):\n",
    "    solver.Add(\n",
    "        solver.Sum([works_shift[(i, j)] for i in range(num_nurses)]) <= 2)\n",
    "  # If s nurses works shifts 2 or 3 on, he must also work that shift the previous\n",
    "  # day or the following day.\n",
    "  solver.Add(\n",
    "      solver.Max(nurses[(2,\n",
    "                         0)] == nurses[(2,\n",
    "                                        1)], nurses[(2,\n",
    "                                                     1)] == nurses[(2,\n",
    "                                                                    2)]) == 1)\n",
    "  solver.Add(\n",
    "      solver.Max(nurses[(2,\n",
    "                         1)] == nurses[(2,\n",
    "                                        2)], nurses[(2,\n",
    "                                                     2)] == nurses[(2,\n",
    "                                                                    3)]) == 1)\n",
    "  solver.Add(\n",
    "      solver.Max(nurses[(2,\n",
    "                         2)] == nurses[(2,\n",
    "                                        3)], nurses[(2,\n",
    "                                                     3)] == nurses[(2,\n",
    "                                                                    4)]) == 1)\n",
    "  solver.Add(\n",
    "      solver.Max(nurses[(2,\n",
    "                         3)] == nurses[(2,\n",
    "                                        4)], nurses[(2,\n",
    "                                                     4)] == nurses[(2,\n",
    "                                                                    5)]) == 1)\n",
    "  solver.Add(\n",
    "      solver.Max(nurses[(2,\n",
    "                         4)] == nurses[(2,\n",
    "                                        5)], nurses[(2,\n",
    "                                                     5)] == nurses[(2,\n",
    "                                                                    6)]) == 1)\n",
    "  solver.Add(\n",
    "      solver.Max(nurses[(2,\n",
    "                         5)] == nurses[(2,\n",
    "                                        6)], nurses[(2,\n",
    "                                                     6)] == nurses[(2,\n",
    "                                                                    0)]) == 1)\n",
    "  solver.Add(\n",
    "      solver.Max(nurses[(2,\n",
    "                         6)] == nurses[(2,\n",
    "                                        0)], nurses[(2,\n",
    "                                                     0)] == nurses[(2,\n",
    "                                                                    1)]) == 1)\n",
    "\n",
    "  solver.Add(\n",
    "      solver.Max(nurses[(3,\n",
    "                         0)] == nurses[(3,\n",
    "                                        1)], nurses[(3,\n",
    "                                                     1)] == nurses[(3,\n",
    "                                                                    2)]) == 1)\n",
    "  solver.Add(\n",
    "      solver.Max(nurses[(3,\n",
    "                         1)] == nurses[(3,\n",
    "                                        2)], nurses[(3,\n",
    "                                                     2)] == nurses[(3,\n",
    "                                                                    3)]) == 1)\n",
    "  solver.Add(\n",
    "      solver.Max(nurses[(3,\n",
    "                         2)] == nurses[(3,\n",
    "                                        3)], nurses[(3,\n",
    "                                                     3)] == nurses[(3,\n",
    "                                                                    4)]) == 1)\n",
    "  solver.Add(\n",
    "      solver.Max(nurses[(3,\n",
    "                         3)] == nurses[(3,\n",
    "                                        4)], nurses[(3,\n",
    "                                                     4)] == nurses[(3,\n",
    "                                                                    5)]) == 1)\n",
    "  solver.Add(\n",
    "      solver.Max(nurses[(3,\n",
    "                         4)] == nurses[(3,\n",
    "                                        5)], nurses[(3,\n",
    "                                                     5)] == nurses[(3,\n",
    "                                                                    6)]) == 1)\n",
    "  solver.Add(\n",
    "      solver.Max(nurses[(3,\n",
    "                         5)] == nurses[(3,\n",
    "                                        6)], nurses[(3,\n",
    "                                                     6)] == nurses[(3,\n",
    "                                                                    0)]) == 1)\n",
    "  solver.Add(\n",
    "      solver.Max(nurses[(3,\n",
    "                         6)] == nurses[(3,\n",
    "                                        0)], nurses[(3,\n",
    "                                                     0)] == nurses[(3,\n",
    "                                                                    1)]) == 1)\n",
    "  # Create the decision builder.\n",
    "  db = solver.Phase(shifts_flat, solver.CHOOSE_FIRST_UNBOUND,\n",
    "                    solver.ASSIGN_MIN_VALUE)\n",
    "  # Create the solution collector.\n",
    "  solution = solver.Assignment()\n",
    "  solution.Add(shifts_flat)\n",
    "  collector = solver.AllSolutionCollector(solution)\n",
    "\n",
    "  solver.Solve(db, [collector])\n",
    "  print(\"Solutions found:\", collector.SolutionCount())\n",
    "  print(\"Time:\", solver.WallTime(), \"ms\")\n",
    "  print()\n",
    "  # Display a few solutions picked at random.\n",
    "  a_few_solutions = [859, 2034, 5091, 7003]\n",
    "\n",
    "  for sol in a_few_solutions:\n",
    "    print(\"Solution number\", sol, \"\\n\")\n",
    "\n",
    "    for i in range(num_days):\n",
    "      print(\"Day\", i)\n",
    "      for j in range(num_nurses):\n",
    "        print(\"Nurse\", j, \"assigned to task\",\n",
    "              collector.Value(sol, shifts[(j, i)]))\n",
    "      print()\n",
    "\n",
    "\n",
    "main()\n",
    "\n"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}
